{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "22d759d0",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/pose.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98b25679",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import mpld3\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from pydrake.all import DiagramBuilder, StartMeshcat\n",
    "\n",
    "from manipulation import FindResource, running_as_notebook\n",
    "from manipulation.scenarios import AddMultibodyTriad\n",
    "from manipulation.station import (\n",
    "    AddPointClouds,\n",
    "    MakeHardwareStation,\n",
    "    load_scenario,\n",
    ")\n",
    "from manipulation.utils import RenderDiagram\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df8331e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd899d5",
   "metadata": {
    "colab_type": "text",
    "id": "7q0A14bAilIX"
   },
   "source": [
    "# Simulating an RGB-D camera\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e3bb7c7",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ILYLouFTjv6e"
   },
   "outputs": [],
   "source": [
    "scenario_data = \"\"\"\n",
    "directives:\n",
    "- add_model:\n",
    "    name: mustard\n",
    "    file: package://drake/manipulation/models/ycb/sdf/006_mustard_bottle.sdf\n",
    "- add_weld:\n",
    "    parent: world\n",
    "    child: mustard::base_link_mustard\n",
    "    X_PC:\n",
    "        translation: [0, 0, 0.09515]\n",
    "        rotation: !Rpy { deg: [-90, 0, -90]}\n",
    "- add_model:\n",
    "    name: camera\n",
    "    file: package://manipulation/camera_box.sdf\n",
    "- add_weld:\n",
    "    parent: world\n",
    "    child: camera::base\n",
    "    X_PC:\n",
    "        translation: [0.5, 0.1, 0.2]\n",
    "        # Point slightly down towards camera\n",
    "        # RollPitchYaw(0, -0.2, 0.2) @ RollPitchYaw(-np.pi/2, 0, np.pi/2)\n",
    "        rotation: !Rpy { deg: [-100, 0, 100] }\n",
    "cameras:\n",
    "    main_camera:\n",
    "        name: camera0\n",
    "        depth: True\n",
    "        X_PB:\n",
    "            base_frame: camera::base\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def DepthCameraDemoSystem():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    scenario = load_scenario(data=scenario_data)\n",
    "    station = builder.AddSystem(MakeHardwareStation(scenario, meshcat))\n",
    "\n",
    "    # Export the camera outputs\n",
    "    builder.ExportOutput(\n",
    "        station.GetOutputPort(\"camera0.rgb_image\"), \"rgb_image\"\n",
    "    )\n",
    "    builder.ExportOutput(\n",
    "        station.GetOutputPort(\"camera0.depth_image\"), \"depth_image\"\n",
    "    )\n",
    "\n",
    "    to_point_cloud = AddPointClouds(\n",
    "        scenario=scenario, station=station, builder=builder, meshcat=meshcat\n",
    "    )\n",
    "\n",
    "    # Add a box for the camera in the environment.\n",
    "    plant = station.GetSubsystemByName(\"plant\")\n",
    "    camera_instance = plant.GetModelInstanceByName(\"camera\")\n",
    "    scene_graph = station.GetSubsystemByName(\"scene_graph\")\n",
    "    AddMultibodyTriad(\n",
    "        plant.GetFrameByName(\"base\", camera_instance),\n",
    "        scene_graph,\n",
    "        length=0.1,\n",
    "        radius=0.005,\n",
    "    )\n",
    "\n",
    "    # Export the point cloud output.\n",
    "    builder.ExportOutput(\n",
    "        to_point_cloud[\"camera0\"].point_cloud_output_port(), \"point_cloud\"\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    diagram.set_name(\"depth_camera_demo_system\")\n",
    "    return diagram\n",
    "\n",
    "\n",
    "def plot_camera_images():\n",
    "    system = DepthCameraDemoSystem()\n",
    "\n",
    "    # Evaluate the camera output ports to get the images.\n",
    "    context = system.CreateDefaultContext()\n",
    "    system.ForcedPublish(context)\n",
    "    color_image = system.GetOutputPort(\"rgb_image\").Eval(context)\n",
    "    depth_image = system.GetOutputPort(\"depth_image\").Eval(context)\n",
    "\n",
    "    # Plot the two images.\n",
    "    plt.subplot(121)\n",
    "    plt.imshow(color_image.data)\n",
    "    plt.title(\"Color image\")\n",
    "    plt.subplot(122)\n",
    "    plt.imshow(np.squeeze(depth_image.data))\n",
    "    plt.title(\"Depth image\")\n",
    "    # mpld3.display()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_camera_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e608aa4",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wya-_6_3MUa1"
   },
   "outputs": [],
   "source": [
    "RenderDiagram(DepthCameraDemoSystem(), max_depth=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baee981d",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mFNDRsZ1MUa4"
   },
   "outputs": [],
   "source": [
    "def plot_manipulation_station_camera_images():\n",
    "    scenario = load_scenario(\n",
    "        filename=FindResource(\"models/clutter.scenarios.yaml\"),\n",
    "        scenario_name=\"Mustard\",\n",
    "    )\n",
    "    station = MakeHardwareStation(scenario, meshcat)\n",
    "\n",
    "    context = station.CreateDefaultContext()\n",
    "\n",
    "    index = 1\n",
    "    plt.figure(figsize=(6, 12))\n",
    "    for camera_num in range(6):\n",
    "        color_image = station.GetOutputPort(\n",
    "            f\"camera{camera_num}.rgb_image\"\n",
    "        ).Eval(context)\n",
    "        depth_image = station.GetOutputPort(\n",
    "            f\"camera{camera_num}.depth_image\"\n",
    "        ).Eval(context)\n",
    "\n",
    "        plt.subplot(6, 2, index)\n",
    "        plt.imshow(color_image.data)\n",
    "        index += 1\n",
    "        plt.title(\"Color image\")\n",
    "        plt.subplot(6, 2, index)\n",
    "        plt.imshow(np.squeeze(depth_image.data))\n",
    "        index += 1\n",
    "        plt.title(\"Depth image\")\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_manipulation_station_camera_images()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
